try: 
    import pyx
    from pyx.graph import *
except ImportError:
    pass

import extensions

import pyx, scipy, operator
from pyx.graph import *

from numpy import *

from unit import G,C
import numpyUtil as nu
import particles as parts
import properties as props
import cosmology, util, ipyx, plot, analytic, library

def talkPlots(qty, cdata, fnbase='talk'):
    args = dict(width=5, backgroundattrs=[pyx.deco.filled([pyx.color.rgb.white])])

    theArgs = util.dictUnion(BoltonSigmaRPlot.defaults, args)
    CosmoSigmaRPlot(cdata).draw(fn=fnbase+'-sigma-r-1', **theArgs)    
    TalkSigmaRPlot(cdata).draw(fn=fnbase+'-sigma-r-2', **theArgs)
    # FIXME
    #SigmaRPlot(cdata).draw(fn=fnbase+'-sigma-r-3', **theArgs)

    theArgs = util.dictUnion(BoltonSigmaMassPlot.defaults, args)
    CosmoSigmaMassPlot(cdata).draw(fn=fnbase+'-sigma-mass-1', **theArgs)    
    TalkSigmaMassPlot(cdata).draw(fn=fnbase+'-sigma-mass-2', **theArgs)
    # FIXME
    #SigmaMassPlot(cdata).draw(fn=fnbase+'-sigma-mass-3', **theArgs)

    args = dict(ratio = 1, width=3, height=3, 
                x=axis.log(min=10**0.2, max=10**1.5, parter=None),
                y=axis.log(min=10**0.2, max=10**1.5, parter=None))
                
    # FIXME
    #LumPlanePlot(qty).draw(fn=fnbase+'-lum-plane', **args)
    #MassPlanePlot(qty).draw(fn=fnbase+'-mass-plane', **args)

    #LensingMassPlot(qty).draw(fn=fnbase+'-lensing-mass', **args)
    
def paperPlots(qty, cdata, multiqty, multicdata, fnbase=''):
    # cdata = [(z, simsAtZ(qty, z)) for z in (0.2, 0.4)]
    RMassPlot(qty, cdata, multiqty, multicdata).draw(fn=fnbase+'f1a', width=6)
    p = plot.Matrix([[SigmaMassPlot(cdata, multicdata)],
                     [SigmaRPlot(cdata, multicdata)]], pdf=True)
    p.draw(fn=fnbase+'f1b', width=6)

    p = plot.Matrix([[LumPlanePlot(qty, multiqty, name="LFP"),
                      MassPlanePlot(qty, multiqty, name="MFP")]], pdf=True)
    p.draw(fn=fnbase+'f2', width=6)

    LensingMassPlot(qty, multiqty).draw(fn=fnbase+'f3', width=6)

def interestingPlots(qty, cosmodata):
    # cdata = [(z, simsAtZ(qty, z)) for z in (0.2, 0.4)]
    lensingSummary(qty, cosmodata)
    lensingESummary(qty, cosmodata)

    LensingPlots().draw()
    LensingPlots(zs = (0.1, 0.3, 0.5, 0.7, 0.9), max=1,
                 fn='lowzLensingPlots').draw()

    EinsteinRadiusOptimalPlot().draw()
    EinsteinRadiusOptimalPlot(zs = (0.1, 0.3, 0.5, 0.7, 0.9), max=1).draw(
        fn='lowzEinsteinRadiusOptimalPlot')


##################################################
# Basic quantities related to lensing.
# These expect  to get a Cosmology instance as the first argument in
# order to calculate distances.

# 10^12 msun => rs = 1e-4 kpc
# 10^10 msun => rs = 1e-6 kpc

def theta_E(cos, zl, zs, rs=1e-5):
    "Einstein radius in radians"
    dl, ds, dls = cos.da(zl), cos.da(zs), cos.da(zs, zl)
    return sqrt(2*rs*dls/ (dl*ds))

def r_E(cos, zl, zs, rs=1e-5):
    "Proper Einstein radius at the lens"
    dl, ds, dls = cos.da(zl), cos.da(zs), cos.da(zs, zl)
    return sqrt(2*rs*dls*dl / ds)

def dtau_dzl(cos, zl, zs, rs=1e-5, n0=1e-9):
    """Given a source at zs, optical depth to lens it with a lens
    between zl and zl+dz, assuming constant comoving number
    density of lenses"""
    return cos.np(zl, n0) * cos.drpdz(zl) * pi * r_E(cos, zl, zs, rs=rs)**2

def dtau_dzs(cos, zl, zs, rs=1e-5, n0=1e-9):
    """Given a lens at zl, optical depth for it to lens a source
    between zs and zs+dz, assuming constant comoving number
    density of sources."""
    return cos.np(zs, n0) * cos.drpdz(zs) * pi * r_E(cos, zl, zs, rs=rs)**2

##################################################
# Stuff for property database
def oldTotalProjectedMassVsRadius(s, p, angles=(0.0, 0.0, 0.0), edges=None):
    if edges is None: edges = arange(0.0, 10.1, 0.5)
    else: edges = asarray(edges)
    
    s.toBodyCoordinates(p)    
    rs = conc((s.rs, s.rg, s.rd))
    ms = conc((s.ms, s.mg, s.md))
    rs = parts.eulerPassive(rs, *angles)

    # Want twice virial radius
    # Find radius correspoding to 200 times cosmic mean density at z=1
    rhov = 200 * 1.48e-8 * (0.71/0.73)**2 * (1+0.0)**3
    rmag = parts.mag(rs)
    order = rmag.argsort()
    densities= ms[order].cumsum() / (4*pi*rmag[order]**3 / 3.0)
    rvir = rmag[nonzero(densities/rhov < 1.0)[0][0]]

    # clip everything outside of 2 rvir
    keep = parts.mag(rs) < rvir
    rs = rs[keep]
    ms = ms[keep]
    
    mBins, junk = nu.histo(parts.mag(rs[:,0:2]), edges, weights=ms, overflow=False)
    return edges[1:], array(mBins).cumsum()

def totalProjectedMassVsRadius(rs, ms, angles=(0.0, 0.0, 0.0), edges=None):
    if edges is None: edges = arange(0.0, 10.1, 0.5)
    else: edges = asarray(edges)

    rs = parts.eulerPassive(rs, *angles)   
    mBins, junk = nu.histo(parts.mag(rs[:,0:2]), edges, weights=ms, overflow=False)
    return edges[1:], array(mBins).cumsum()

def totalMassVsRadius(rs, ms, angles=(0.0, 0.0, 0.0), edges=None):
    if edges is None: edges = arange(0.0, 10.1, 0.5)
    else: edges = asarray(edges)

    rs = parts.eulerPassive(rs, *angles)   
    mBins, junk = nu.histo(parts.mag(rs), edges, weights=ms, overflow=False)
    return edges[1:], array(mBins).cumsum()

##################################################
## Plots        
class EinsteinAnglePlot (plot.Cosmology):        
    defaults = dict(key=key.key(pos='tr'),
                    x=axis.lin(title='$z$ Lens'),
                    y=axis.slog(min=0.1, max=2,title=r'$\Theta_E$ (arcsec)'))

    def data(self, cosmo):
        return [data.function('y(x)=206265*theta_E(cosmo, x, z)',
                              min=0.02, max=z, title=None,
                              context=dict(theta_E=theta_E, 
                                           z=z, cosmo=cosmo))
                for z in self.zs]

class EinsteinAngleSourcePlot (plot.Cosmology):        
    defaults = dict(key=key.key(pos='tl'),
                    x=axis.lin(title='$z$ Source'),
                    y=axis.slog(min=0.3, max=6,title=r'$\Theta_E$ (arcsec)'))

    def data(self, cosmo):
        return [data.function('y(x)=206265*theta_E(cosmo, z, x)',
                              min=z, max=self.max, title=None, 
                              context=dict(theta_E=theta_E,
                                           z=z, cosmo=cosmo))
                for z in self.zs]

class EinsteinRadiusPlot (plot.Cosmology):
    defaults = dict( key=key.key(pos='tr'),
                     x=axis.lin(title='$z$ Lens'),
                     y=axis.lin(title=r'$R_{\rm Ein}$ (kpc)'))
    
    def data(self, cosmo):
        return [data.function('y(x)=r_E(cosmo, x, z)',
                              min=0.02, max=z, title=None, 
                              context=dict(r_E=r_E,
                                           z=z, cosmo=cosmo))
                for z in self.zs]

class EinsteinRadiusSourcePlot(plot.Cosmology):
    defaults = dict(key=key.key(pos='tl'),
                    x=axis.lin(title='$z$ Source'),
                    y=axis.lin(title=r'$R_{\rm Ein}$ (kpc)'))

    def data(self, cosmo):
        return [data.function('y(x)=r_E(cosmo, z, x)',
                              min=z, max=self.max, title=None, 
                              context=dict(r_E=r_E,
                                           z=z, cosmo=cosmo))
                for z in self.zs]

class OpticalDepthPlot(plot.Cosmology):
    defaults = dict(key=key.key(pos='tr'),
                    x=axis.lin(title='$z$ Lens'),
                    y=axis.lin(title=r'$d\tau/dz$'))
        
    def data(self, cosmo):
        return [data.function('y(x)=dtau_dzl(cosmo, x, z)',
                              min=0.02, max=z, title=None,
                              context=dict(dtau_dzl=dtau_dzl,
                                           z=z, cosmo=cosmo))
                for z in self.zs]
    
class OpticalDepthSourcePlot (plot.Cosmology):
    defaults = dict(key=key.key(pos='tl'),
                    x=axis.lin(title='$z$ Source'),
                    y=axis.lin(title=r'$d\tau/dz$'))

    def data(self, cosmo):
        return [data.function('y(x)=dtau_dzs(cosmo, z, x)',
                              min=z, max=self.max, title=None,
                              context=dict(dtau_dzs=dtau_dzs,
                                           z=z, cosmo=cosmo))
                for z in self.zs]
    
class EinsteinRadiusOptimalPlot (plot.Plot):
    defaults = dict(key=key.key(pos='tr'),
                    x=axis.lin(title='$z$ Lens'),
                    y=axis.lin(title=r'$R_{\rm Ein}$ (kpc)', max=17.5))
    
    def __init__(self, max=5, **kw):
        self.max = max

    def data(self):
        cosmo = cosmology.Standard()
        return [data.function('y(x)=r_E(cosmo, x, 2*x, rs=1e-6*fmass)',
                              title=r'$%d\times10^{10} M_\odot$' % fmass,
                              min=0.02, max=self.max,
                              context=dict(r_E=r_E,
                                           fmass=fmass, cosmo=cosmo))
                for fmass in (1, 3, 10, 30, 100)]
                
class LensingPlots (plot.Matrix):
    defaults = dict(key=None)
    
    def __init__(self, **kw):
        kw['key'] = None
        arr = [ [EinsteinAnglePlot(**kw), EinsteinAngleSourcePlot(**kw)],
                [EinsteinRadiusPlot(**kw), EinsteinRadiusSourcePlot(**kw)],
                [OpticalDepthPlot(**kw), OpticalDepthSourcePlot(**kw)] ]
        plot.Matrix.__init__(self, arr)

    def data(self):
        result = plot.Matrix.data(self)
        result[-1][-1].key = key.key(pos='tl')
        return result

##################################################
## Plot just the observational data

class BoltonSigmaRPlot(plot.Plot):
    defaults = dict(x=axis.log(min=170, max=410, title=r'$\sigma_{e2}$ (km/s)'),
                    y=axis.log(min=1, max=10, title=r'$R_{\rm Ein}$ (kpc)'))

    def dataBolton(self):
        return data.points(scipy.io.read_array('bolton.txt', columns=(1,3)), x=1, y=2)

class BoltonSigmaMassPlot(plot.Plot):
    defaults = dict(x=axis.log(min=170, max=410, title=r'$\sigma_{e2}$ (km/s)'),
                    y=axis.log(min=2e10, max=1e12, title=r'$M_{\rm Ein} (M_\odot)$'))

    def dataBolton(self):
        return data.points(scipy.io.read_array('bolton.txt', columns=(1,4)), x=1, y=2)

class BoltonRMassPlot(plot.Plot):
    defaults = defaults = dict(x=axis.log(min=1, max=10, title=r'$R_{\rm Ein}$ (kpc)'),
                               y=axis.log(min=2e10, max=1e12, title=r'$M_{\rm Ein} (M_\odot)$'))
    
    def dataBolton(self):
        return data.points(scipy.io.read_array('bolton.txt', columns=(3,4)), x=1, y=2)
        
##################################################
## Plot continuous quantities from simulations
class SimPlot(plot.Plot):
    def __init__(self, qty, multiqty=None):        
        self.qty = qty
        if multiqty is None: self.multiqty = {}
        else:                self.multiqty = multiqty

    def plotSims(self):
        self.g.plot(self.dataSims(self.qty), [style.line([pyx.style.linestyle.solid,
                                                          pyx.color.rgb.blue,
                                                          pyx.color.transparency(0.9)])])
        self.g.plot(self.dataSims(self.multiqty), [style.line([pyx.style.linestyle.solid,
                                                          pyx.color.rgb.red,
                                                          pyx.color.transparency(0.1)])])
    def finish (self, fn): 
        self.g.writePDFfile(fn)

class SimRMassPlot(SimPlot):
    defaults = dict(x=axis.log(min=0.5, title=r'$R_{\rm Ein}$ (kpc)'),
                    y=axis.log(min=2e8, max=4e11, title=r'$M_{\rm Ein} (M_\odot)$'))

    def dataSims(self, qty):
        return [data.array([xs, 1e10*ys], x=1, y=2)
                for xs, ys in props.mapSnap(props.fetcher('surface-mass-off-r',
                                                          'surface-mass-off'), qty)]
##################################################
## These show velocity dispersion within einstein radius rather than
## half mass radius.  But the half mass radius is independent of
## redshift, so the vel disp w/in half mass radius in indep of
## redshift too
class SimSigmaRPlot(SimPlot):
    defaults = dict(x=axis.log(min=20,max=400, title=r'$\sigma$ (km/s)'),
                    y=axis.log(min=0.5, title=r'$R_{\rm Ein}$ (kpc)'))
    
    def dataSims(self, qty):
        return [data.array([xs, ys], x=1, y=2)
                for xs, ys in props.mapSnap(props.fetcher('aperture-sigma-off',
                                                          'aperture-sigma-off-r'), qty)]
class SimSigmaMassPlot(SimPlot):
    defaults = dict(x=axis.log(min=20, max=400, title=r'$\sigma$ (km/s)'),
                    y=axis.log(min=2e8, max=4e11, title=r'$M_{\rm Ein} (M_\odot)$'))

    def dataSims(self, qty):
        return [data.array([xs, 1e10*ys], x=1, y=2)
                for xs, ys in props.mapSnap(props.fetcher('aperture-sigma-off',
                                                          'surface-mass-off'), qty)]

##################################################
## Put simulations at specific redshifts
def simsAtZ(qty, z):
    c = cosmology.Standard()

    def eqn(x, z, mr):
        "Free variables: c"
        return r_E(c, z, 2*z, rs=2*G*mr(x)[0]/C**2) - x

    result = []
    for mRs, ms, sRs, ss, reff, sigma2, sim, snap in \
        props.mapSnap(props.fetcher('surface-mass-off-r',
                                    'surface-mass-off',
                                    'aperture-sigma-off-r',
                                    'aperture-sigma-off',
                                    'reff-off', 'sigma-off-2',
                                    'simulation', 'snapshot'), qty):
        mr = scipy.interpolate.interp1d(mRs,ms)
        sr = scipy.interpolate.interp1d(sRs,ss)
        if eqn(0.5, z, mr) < 0:
            continue # Not lensed
        if eqn(mRs[-1], z, mr) > 0:
            print "Warning, %s %s has large einstein radius!" % (sim, snap)
            continue # Einstein radius large..

        rE = scipy.optimize.bisect(lambda x: eqn(x, z, mr), 0.5, mRs[-1])
        result.append((rE, 1e10*mr(rE)[0], sr(rE)[0], reff, sigma2))
    return data.points(result, r=1, m=2, sigma=3, reff=4, sigma2=5)

class SimCosmoPlot(plot.Plot):
    def __init__(self, cosmoData, multiData=None):
        """Cosmodata is a list of tuples where the first element is a
        redshift and the second eleemnt is a Pyx data object with
        columns r, m, and sigma"""
        self.cosmoData = cosmoData
        if multiData: self.multiData = multiData
        else:         self.multiData = []

    def plotCosmo(self):
        for z, dat in self.cosmoData:
            self.g.plot(self.dataCosmo(dat),
                        [style.symbol(style.symbol.circle, size=0.3*z,
                                      symbolattrs=[pyx.deco.filled,
                                                   pyx.color.rgb.blue,
                                                   pyx.color.transparency(0.25)])])
        for z, dat in self.multiData:
            self.g.plot(self.dataCosmo(dat),
                        [style.symbol(style.symbol.circle, size=0.3*z,
                                      symbolattrs=[pyx.deco.filled,
                                                   pyx.color.rgb.red,
                                                   pyx.color.transparency(0.25)])])
        
    def finish (self, fn): 
        self.g.writePDFfile(fn)

class CosmoSigmaRPlot(SimCosmoPlot):
    defaults = dict(x=axis.log(title=r'$\sigma$ (km/s)'),
                    y=axis.log(title=r'$R_{\rm Ein}$ (kpc)'))

    def dataCosmo(self, dat):
        return data.data(dat, x='sigma2', y='r')
        
class CosmoSigmaMassPlot(SimCosmoPlot):
    defaults = dict(x=axis.log(title=r'$\sigma$ (km/s)'),
                    y=axis.log(title=r'$M_{\rm Ein} (M_\odot)$'))

    def dataCosmo(self, dat):
        return data.data(dat, x='sigma2', y='m')

class CosmoRMassPlot(SimCosmoPlot):
    defaults = dict(x=axis.log(title=r'$R_{\rm Ein}$ (kpc)'),
                    y=axis.log(title=r'$M_{\rm Ein} (M_\odot)$'))
    
    def dataCosmo(self, dat):
        return data.data(dat, x='r', y='m')

##################################################
## These show velocity dispersion within einstein radius rather than
## half mass radius.  But the half mass radius is independent of
## redshift, so the vel disp w/in half mass radius in indep of
## redshift too
class CosmoSigmaERPlot(SimCosmoPlot):
    defaults = dict(x=axis.log(title=r'$\sigma$ (km/s)'),
                    y=axis.log(title=r'$R_{\rm Ein}$ (kpc)'))

    def dataCosmo(self, dat):
        return data.data(dat, x='sigma', y='r')
        
class CosmoSigmaEMassPlot(SimCosmoPlot):
    defaults = dict(x=axis.log(title=r'$\sigma$ (km/s)'),
                    y=axis.log(title=r'$M_{\rm Ein} (M_\odot)$'))

    def dataCosmo(self, dat):
        return data.data(dat, x='sigma', y='m')

##################################################
## Analytic Models
def model(m, z):
    """takes mass model and redshift, returns einstein readius."""
    c = cosmology.Standard()
    dl, ds, dls = c.da(z), c.da(2*z), c.da(2*z, z)
    const = 4*G*dls*dl/(C**2*ds)
    rmin, rmax = 0.01*m._rScale, 95*m._rScale
    theF = lambda x: const*m.fL(x) - x**2
    if theF(rmax)*theF(rmin) > 0:
        if theF(rmin) < 0: print "Rmin is too big"
        if theF(rmax) > 0: print "Rmax is too small"
        return 0.001
    result = scipy.optimize.brentq(theF, rmin, rmax)
    return result

def _slowHernquistModel(mtot, rvir, z):
    m = analytic.Galaxy(rvir=rvir, mtot=mtot).baryons
    rE = model(m, z)
    reffProj = m.reProj()
    return 1e10*m.fL(rE), rE, reffProj, m.fsigma_a(reffProj/2.0)
hernquistModel = util.memoize(_slowHernquistModel)

def _slowNfwModel(mtot, rvir, z):
    m = analytic.Galaxy(rvir=rvir, mtot=mtot).dark
    rE = model(m, z)
    reff = scipy.optimize.brentq(lambda x: m.fL(x) - mtot/2.0, 0, rvir) 
    return 1e10*m.fL(rE), rE, reff, m.fsigma_a(reff/2.0)
nfwModel = util.memoize(_slowNfwModel)

def _slowGalaxyModel(mtot, rvir, z):
    """Returns """
    m = analytic.Galaxy(rvir=rvir, mtot=mtot)
    rE = model(m, z)
    # mE, rE refer to total mass.  Reff refers to light.  sigma
    # _responds_ to total mass, but _measured_ at 1/2 of half light
    # radius
    reffProj = m.baryons.reProj()
    return 1e10*m.fL(rE), rE, reffProj, m.fsigma_a(reffProj/2.0)
galaxyModel = util.memoize(_slowGalaxyModel)

class ModelPlot(plot.Plot):
#     def plotHernquist(self): self.modelPlot(self.dataHernquist, pyx.style.linestyle.dotted)
#     def plotNfw(self): self.modelPlot(self.dataNfw, pyx.style.linestyle.dashed)
    def plotGalaxy(self): self.modelPlot(self.dataGalaxy, pyx.style.linestyle.solid)
        
    def dataGalaxy(self, *a): return self.theData(galaxyModel, *a)
#     def dataNfw(self, *a): return self.theData(nfwModel, *a)
#     def dataHernquist(self, *a): return self.theData(hernquistModel, *a)
        
    def modelPlot(self, datafn, exls):
        mvirs = (10, 30, 100, 300, 1000)
        rvirs = (25, 50, 100, 200, 400, 800)
        
        for z, ls in ((0.2, pyx.style.linestyle.solid),):
                      # (0.4, pyx.style.linestyle.dotted)):

            theSty = [style.line([pyx.color.rgb.green,
                                  pyx.color.transparency(0.5),
                                  exls, ls]),
                      style.symbol(symbol=style.symbol.circle, size=0.05,
                                   symbolattrs=[pyx.color.rgb.green, exls,
                                                pyx.color.transparency(0.5)])]
            # A separate line for each m is better, if the grid is too
            # busy
            for m in mvirs:
                theDat = datafn(z, [m], rvirs)
                self.g.plot(theDat, theSty)
            for r in rvirs:
                theDat = datafn(z, mvirs, [r])
                self.g.plot(theDat, theSty)
            
class SisSigmaMassPlot (plot.Plot):
    def plotSis(self):
        for z, ls in ((0.2, pyx.style.linewidth.normal),):
                      #(0.4, pyx.style.linewidth.Thin)):
            self.g.plot(self.dataSis(z),
                        [style.line([ls, pyx.color.rgb.green])])

    def dataSis(self, z):
        c = cosmology.Standard()
        dl, ds, dls = c.da(z), c.da(2*z), c.da(2*z, z)            
        return data.function('y(x) = 1e10*const*x**4',
                             context=dict(const=4*pi**2*dl*dls/(G*C**2*ds)))
    
class ModelSigmaMassPlot(ModelPlot):
    defaults = dict(x=axis.log(title=r'$\sigma$ (km/s)', min=40,max=400),
                    y=axis.log(title=r'$M_{\rm Ein} (M_\odot)$', min=1e9, max=1e12))

    def theData(self, f, z, mvirs, rvirs):        
        theDat = [f(mvir, rvir, z)
                  for mvir, rvir in util.crossSet(mvirs, rvirs)
                  if f(mvir, rvir, z)[1] > 0.3]
        if theDat: return data.points(theDat, x=4, y=1)
        else:      return data.points([[None, None]], x=1, y=2)

    def finish (self, fn): 
        self.g.writePDFfile(fn)


class ModelSigmaRPlot(ModelPlot):
    defaults = dict(x=axis.log(title=r'$\sigma$ (km/s)', min=40,max=400),
                    y=axis.log(title=r'$R_{\rm Ein}$ (kpc)', min=1,max=30))

    def theData(self, f, z, mvirs, rvirs):
        theDat = [f(mvir, rvir, z) 
                  for mvir, rvir in util.crossSet(mvirs, rvirs)
                  if f(mvir, rvir, z)[1] > 0.3]
        if theDat: return data.points(theDat, x=4, y=2)
        else:      return data.points([[None, None]], x=1, y=2)

class ModelRMassPlot(ModelPlot):
    defaults = dict(x=axis.log(title=r'$R_{\rm Ein}$ (kpc)',min=1,max=30),
                    y=axis.log(title=r'$M_{\rm Ein} (M_\odot)$', min=1e9,max=1e12))

    def theData(self, f, z, mvirs, rvirs):
        theDat = [f(mvir, rvir, z) 
                  for mvir, rvir in util.crossSet(mvirs, rvirs)]
        if theDat: return data.points(theDat, x=2, y=1)
        else:      return data.points([[None, None]], x=1, y=2)

##################################################
## Intermediate stages for build-up during talk
class TalkSigmaMassPlot (BoltonSigmaMassPlot, CosmoSigmaMassPlot):
    def __init__(self, cosmodata):
        CosmoSigmaMassPlot.__init__(self, cosmodata)
        
    def finish (self, fn): 
        self.g.writePDFfile(fn)
    
class TalkSigmaRPlot (BoltonSigmaRPlot, CosmoSigmaRPlot):
    def __init__(self, cosmodata):
        CosmoSigmaRPlot.__init__(self, cosmodata)    

    def finish (self, fn): 
        self.g.writePDFfile(fn)

##################################################
## Everything
class SigmaRPlot (BoltonSigmaRPlot, CosmoSigmaRPlot, ModelSigmaRPlot):
    def __init__(self, cosmodata, mdata=None):
        CosmoSigmaRPlot.__init__(self, cosmodata, mdata)    

    def finish (self, fn): 
        self.g.writePDFfile(fn)

# class SigmaMassPlot (BoltonSigmaMassPlot, CosmoSigmaMassPlot,
#                      ModelSigmaMassPlot, SisSigmaMassPlot):
class SigmaMassPlot (BoltonSigmaMassPlot, CosmoSigmaMassPlot, ModelSigmaMassPlot):
    def __init__(self, cosmodata, multidata=None):
        CosmoSigmaMassPlot.__init__(self, cosmodata, multidata)    
        
    def finish (self, fn): 
        self.g.writePDFfile(fn)

class RMassPlot (BoltonRMassPlot, SimRMassPlot, CosmoRMassPlot):
    def __init__(self, qty, cosmodata, multiqty=None, multidata=None):
        SimRMassPlot.__init__(self, qty, multiqty)
        CosmoRMassPlot.__init__(self, cosmodata, multidata)

    def finish (self, fn): 
        self.g.writePDFfile(fn)

class SigmaERPlot (BoltonSigmaRPlot, SimSigmaRPlot, CosmoSigmaERPlot, ModelRMassPlot):
    def __init__(self, qty, cosmodata):
        SimSigmaRPlot.__init__(self, qty)
        CosmoSigmaERPlot.__init__(self, cosmodata)    

    def finish (self, fn): 
        self.g.writePDFfile(fn)

class SigmaEMassPlot (BoltonSigmaMassPlot, SimSigmaMassPlot,
                      CosmoSigmaEMassPlot, ModelSigmaMassPlot,
                      SisSigmaMassPlot):
    def __init__(self, qty, cosmodata, **kw):
        SimSigmaMassPlot.__init__(self, qty)
        CosmoSigmaEMassPlot.__init__(self, cosmodata)    
        ModelSigmaMassPlot.__init__(self, **kw)

    def finish (self, fn): 
        self.g.writePDFfile(fn)

def lensingSummary(qty, cosmodata, fn='lensingSummary'):
    c = pyx.canvas.canvas()
    
    for plot, ypos in ((RMassPlot(qty, cosmodata), 16),
                       (SigmaMassPlot(cosmodata), 8),
                       (SigmaRPlot(cosmodata), 0)):
        plot.draw(fn=None, ypos=ypos)
        c.insert(plot.g)

    c.writePDFfile(fn)

def lensingESummary(qty, cosmodata, fn='lensingESummary'):
    c = pyx.canvas.canvas()
    
    for plot, ypos in ((RMassPlot(qty, cosmodata), 16),
                       (SigmaEMassPlot(qty, cosmodata), 8),
                       (SigmaERPlot(qty, cosmodata), 0)):
        plot.draw(fn=None, ypos=ypos)
        c.insert(plot.g)

    c.writePDFfile(fn)

def kappa(qty):
    k1s, k2s, k3s = [], [], []
    for sim, snap, reff, sig, ie in props.mapSnap(props.fetcher('simulation', 'snapshot',
                                                                'reff-off',
                                                    'sigma-off-2', 'mass-surface-density-off-2'),qty):
        k1s.append(log10(sig**2*reff)/sqrt(2))
        k2s.append(log10(sig**2*ie**2/reff)/sqrt(6))
        k3s.append(log10(sig**2/(ie*reff))/sqrt(3))

    return k1s, k2s, k3s
                   

class GsnRSigmaPlot(plot.Plot):
    defaults = dict(x=axis.log(min=100, max=400, title='$\sigma$ (km/s)'),
                    y=axis.log(min=1, max=15, title=r"$R_e$ (kpc)"))

    def __init__(self, qty, multiqty=None,
                 tags=('Sbc', 'GEqualMassMajor', 'cos'), name=None):
        if multiqty: self.multiqty = multiqty
        else: self.multiqty = {}
        
        self.qty, self.tags, self.name = qty, tags, name

    def finish (self, fn): 
        """Declares plot done and writes it"""
        self.g.writePDFfile(fn)

    def plot(self):

        self.g.plot(self.data(self.qty), [style.symbol(style.symbol.circle, size=0.1,
                                 symbolattrs=[pyx.deco.filled,
                                              pyx.color.rgb.blue,
                                              pyx.color.transparency(0.5)])])
        self.g.plot(self.data(self.multiqty), [style.symbol(style.symbol.circle, size=0.1,
                                 symbolattrs=[pyx.deco.filled,
                                              pyx.color.rgb.red,
                                              pyx.color.transparency(0.1)])])        

        return self.g
    
    def data(self, qty):
        allTags = reduce(operator.add, [library.tags[tag]
                                        for tag in self.tags])
        dat = props.mapSnap(props.fetcher('simulation', 'sigma-off-2', 'reff-off'), qty)
        dat = [el[1:] for el in dat if el[0] in allTags]
                 
        return data.points(dat, x=1, y=2)

class KappaPlot(plot.Plot):
    def __init__(self, qty, multiqty=None,
                 tags=('Sbc', 'GEqualMassMajor', 'cos'), name=None):
        if multiqty: self.multiqty = multiqty
        else: self.multiqty = {}
        
        self.qty, self.tags, self.name = qty, tags, name

    def finish (self, fn): 
        """Declares plot done and writes it"""
        self.g.writePDFfile(fn)

    def setup(self, **kw):
        self.kw = kw
        self.g = pyx.canvas.canvas()

    def plot(self):
        g1 = graphxy(height=10, width=10,
                     x = axis.lin(min=1, max=5,
                             title=r'$\kappa_1 = \log (R_e \sigma^2) / \sqrt{2} \propto M$'),
                     y = axis.lin(min=-2, max=2,
          title=r'$\kappa_2 = \log (\sigma^2 I_e^2 / R_e) / \sqrt{6}  \propto (M/L) I_e^3$'))
        g2 = graphxy(height=2.5, width=10, ypos=g1.height + 0.5,
                           y=axis.lin(min=2.25, max=3.25,
                       title=r'$\kappa_3 = \log(\sigma^2/I_e R_e)/\sqrt {3} \propto M/L$'),
                           x=axis.linkedaxis(g1.axes['x']))

        binStyle = [style.symbol(style.symbol.circle, size=0.1,
                                 symbolattrs=[pyx.deco.filled,
                                              pyx.color.rgb.blue,
                                              pyx.color.transparency(0.5)])]        
        cosStyle = [style.symbol(style.symbol.circle, size=0.1,
                                 symbolattrs=[pyx.deco.filled,
                                              pyx.color.rgb.red,
                                              pyx.color.transparency(0.1)])]        

        k1, k2, k3 = self.data(self.qty)
        g1.plot(data.array([k1, k2], x=1, y=2), binStyle)
        g2.plot(data.array([k1, k3], x=1, y=2), binStyle)

        k1, k2, k3 = self.data(self.multiqty)
        g1.plot(data.array([k1, k2], x=1, y=2),cosStyle)
        g2.plot(data.array([k1, k3], x=1, y=2),cosStyle)


        self.g.insert(g1)
        self.g.insert(g2)
        return self.g
    
    def data(self, qty):
        return kappa(qty)
        
class PlanePlot(plot.Plot):
    defaults = defaults = dict(x=axis.log(), y=axis.log())
    
    def __init__(self, qty, multiqty=None,
                 tags=('Sbc', 'GEqualMassMajor', 'cos'), name=None):
        if multiqty: self.multiqty = multiqty
        else: self.multiqty = {}
        
        self.qty, self.tags, self.name = qty, tags, name

    def plotSim(self):
        self.g.plot(self.dataSim(self.qty),
                    [style.symbol(style.symbol.circle, size=0.1,
                                  symbolattrs=[pyx.deco.filled,
                                               pyx.color.rgb.blue,
                                               pyx.color.transparency(0.5)])])

        if self.multiqty:
            self.g.plot(self.dataSim(self.multiqty),
                        [style.symbol(style.symbol.circle, size=0.1,
                                      symbolattrs=[pyx.deco.filled,
                                                   pyx.color.rgb.red,
                                                   pyx.color.transparency(0.1)])])
        if self.name:
            plot.pyxTitle(self.g, self.name, 0.2, 0.8)
            
    def dataSim(self, qty):
        dat = [el for el in props.mapSnap(props.fetcher('simulation',
                                        'sigma-off-2',
                                        'mass-surface-density-off-2', 
                                        'reff-off'), qty) if el]
        if self.tags:
            allTags = reduce(operator.add, [library.tags[tag]
                                            for tag in self.tags])
            dat = [el for el in dat if el[0] in allTags]

        names = [el[0] for el in dat]
        dat = [el[1:] for el in dat]

        sigma2, Sigma2, reff = transpose(dat)         
        Sigma2 = 1e10*Sigma2

        a, b, c = self.a, self.b, self.c
        reffPlane = sigma2**a * Sigma2**b * 10**c
        return data.array([reffPlane, reff], x=1, y=2)

    def plotBolton(self):
        sty1 = [style.line([pyx.style.linestyle.solid])]
        sty2 = [style.line([pyx.style.linestyle.dashed])]
        self.g.plot(self.dataBolton(offset=0.0), sty1)
        self.g.plot(self.dataBolton(offset=0.2), sty2)
        self.g.plot(self.dataBolton(offset=-0.2), sty2)

    def dataBolton(self, **kw):
        return data.function('y(x) = 10**offset * x',
                             title=None, context=kw)
    def finish (self, fn): 
        """Declares plot done and writes it"""
        self.g.writePDFfile(fn)
                       
class LumPlanePlot(PlanePlot):
    a, b, c = 1.5, -0.78, 3.9
    defaults = defaults = dict(ratio = 1,
                               x=axis.log(min=10**0.2, max=10**1.5,
                                          title=r'$10^{3.9} \sigma_{e2}^{1.5} / I_{e2}^{0.8}$'),
                               y=axis.log(min=10**0.2, max=10**1.5,
                                          title=r'$R_e$ (kpc)'))
    
    def dataSim(self, *a, **kw):
        dat = PlanePlot.dataSim(self, *a, **kw)
        # FIXME FIXME FIXME -- why is this here
        dat.columns['x'] = [2.5*el for el in dat.columns['x']]
        return dat
        
class MassPlanePlot(PlanePlot):
    a, b, c = 1.77, -1.16, 7.8

    defaults = defaults = dict(ratio = 1,
                               x=axis.log(min=10**0.2, max=10**1.5,
                                          title=r'$10^{7.8} \sigma_{e2}^{1.8} / \Sigma_{e2}^{1.2}$'),
                               y=axis.log(min=10**0.2, max=10**1.5,
                                          title=r'$R_e$ (kpc)'))

    def dataSim(self, *a, **kw):
        dat = PlanePlot.dataSim(self, *a, **kw)
        # FIXME FIXME FIXME -- why is this here
        dat.columns['x'] = [el/5.5 for el in dat.columns['x']]
        return dat

class LensingMassPlot(plot.Plot):
    defaults = defaults = dict(ratio = 1,
                               x=axis.log(min=10**9.8, max=10**11.7,
                                          title=r'$\sigma_{e2}^2 R_e / 2 G$'),
                               y=axis.log(min=10**10.3, max=10**12.2, title='$M_{e2}$'))

    def __init__(self, qty, multiqty=None, tags=('Sbc', 'GEqualMassMajor', 'cos')):
        self.qty, self.tags = qty, tags
        if multiqty: self.multiqty = multiqty
        else: self.multiqty = {}
        
    def plotBolton(self):
        self.g.plot(self.dataBolton(slope=0.986, offset=0.58))
        # NOTE -- my estimate of Bolton's uncertainty in the offset without covariance:
        # 0.1 dex / sqrt(34)
        doff = 0.1/sqrt(36)

        for s, o, mn, mx in ((0.986 + 0.034, -0.034*10.7 + 0.58 + doff, 5e10, None),
                             (0.986 + 0.034, -0.034*10.7 + 0.58 - doff, None, 5e10),
                             (0.986 - 0.034,0.034*10.7 + 0.58 + doff, None, 5e10),
                             (0.986 - 0.034,0.034*10.7 + 0.58 - doff, 5e10, None)):
            self.g.plot(self.dataBolton(slope=s, offset=o, min=mn, max=mx),
                        [style.line([pyx.style.linestyle.dotted])])

    def dataBolton(self, min=None, max=None, **kw):
        return data.function('y(x) = 10**offset * x**slope',
                             min=min, max=max, title=None, context=kw)
    
    def plotSim(self):
        self.g.plot(self.dataSim(self.qty),
                    [style.symbol(style.symbol.circle, size=0.1,
                                  symbolattrs=[pyx.deco.filled,
                                               pyx.color.rgb.blue,
                                               pyx.color.transparency(0.5)])])

        self.g.plot(self.dataSim(self.multiqty),
                    [style.symbol(style.symbol.circle, size=0.1,
                                  symbolattrs=[pyx.deco.filled,
                                               pyx.color.rgb.red,
                                               pyx.color.transparency(0.1)])])

    def dataSim(self,qty):
        if not qty: return data.points([None, None], x=1, y=2)
        dat = [el for el in props.mapSnap(props.fetcher('simulation',
                                        'sigma-off-2',
                                        'mass-surface-density-off-2', 
                                        'reff-off'), qty) if el]

        if self.tags:
            allTags = reduce(operator.add, [library.tags[tag]
                                            for tag in self.tags])
            dat = [el for el in dat if el[0] in allTags]

        dat = [el[1:] for el in dat]

        sigma2, Sigma2, reff = transpose(dat)         
        xs = 1e10 * sigma2**2 * reff / (2*G)
        ys = 1e10 * Sigma2 * (reff/2)**2
        ys = 10*ys  # FIXME FIXME FIXME Why is this here?
        return data.array([xs, ys], x=1, y=2)

    def finish (self, fn): 
        """Declares plot done and writes it"""
        self.g.writePDFfile(fn)


##################################################
## Rickety code
def dmImportanceDat(qty):    
    rawDat = props.mapSnap(props.fetcher('3-projected-mass-off',
                              '3-mass-off', 'sigma-off-2', 'simulation'), qty)
    dat = zip(*[(a2['stars'][0], a2['dark'][0],
                 a3['stars'][0], a3['dark'][0], name)
                for a2, a3, sig, name in rawDat
                if sig > 180])
    return array(dat[0]), array(dat[1]), array(dat[2]), array(dat[3]), dat[4]

#     ph(d3/(s3 + d3), (0, 1, 20), filename='gsn')
#     ph(d2/(s2 + d2), (0, 1, 20), filename='gsn')

    def finish (self, fn): 
        """Declares plot done and writes it"""
        self.g.writePDFfile(fn)

def summary(fn='lensingSummary'):
    c = pyx.canvas.canvas()
    
    for plot, ypos in ((ModelSigmaMassPlot(), 8),
                       (ModelSigmaRPlot(), 0)):
        plot.draw(fn=None, ypos=ypos, x=axis.tlog(), y=axis.tlog())
        c.insert(plot.g)

    c.writeEPSfile(fn)
